import os
import glob

# 设置路径A
path_A = r'C:\Users\Caibin\Desktop\mixup_20250124_fix\mixup\train\labels'  # 替换为你的txt文件夹路径


def process_file(file_path):
    with open(file_path, 'r') as f:
        lines = f.readlines()

    new_lines = []
    for line in lines:
        parts = line.strip().split()

        # 确保至少有 23 个元素（类别 + 4 个框参数 + 6 对关键点 + 6 个可见性标记）
        if len(parts) < 23:
            continue

        # 保留 类别, center_x, center_y, width, height
        new_line = parts[:5]

        # 从第6个开始都是关键点信息
        keypoints = parts[5:]

        # 提取 6 对关键点 (x, y, visible)
        points = []
        for i in range(0, len(keypoints), 3):
            points.append((keypoints[i], keypoints[i + 1], keypoints[i + 2]))

        if int(parts[0]) == 0:
            # 第五号点移动至第四号点位置
            points[3] = points[4]

            # 第五、六号点修改为 0.000000 0.000000 0
            points[4] = ("0.000000", "0.000000", "0")
            points[5] = ("0.000000", "0.000000", "0")

        # 拼接新的行内容
        for px, py, vis in points:
            new_line.append(px)
            new_line.append(py)
            new_line.append(vis)

        # 组合并加入新行
        new_lines.append(" ".join(new_line))


    # 覆盖写入原文件
    with open(file_path, 'w') as f:
        f.write("\n".join(new_lines))


def main():
    # 查找路径A下的所有txt文件
    txt_files = glob.glob(os.path.join(path_A, "*.txt"))

    # 对每个txt文件进行处理
    for txt_file in txt_files:
        process_file(txt_file)
        print(f"已处理: {txt_file}")


if __name__ == "__main__":
    main()